# ==============================================================================
# 文件: main.py
# 作用: 项目的主入口文件。它作为调度器，调用其他模块的功能，按顺序
#      执行整个PU学习流程以及所有配置的基准对比实验。
# ==============================================================================
import os
import torch
import shutil

# --- 全局设置Hugging Face镜像 ---
# 确保所有从Hugging Face下载的操作都通过国内镜像进行，避免网络问题。
# 这个设置必须在导入transformers相关模块之前完成。
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# 导入项目内的其他模块
import config
import data_loader
import feature_extractor
import initial_trainer
import em_trainer
import baselines
import evaluator
from models import SimpleMLP
from torch.utils.data import TensorDataset

def main():
    """主执行函数"""
    print("="*60)
    print("开始执行PU学习及基准对比流程")
    print(f"使用设备: {config.DEVICE}")
    print("="*60)

    # --- 自动清理并重建模型目录，确保每次都重新训练模型 ---
    # config中定义了模型文件的完整路径，.parent可以获取其父目录
    initial_model_dir = config.INITIAL_MODEL_PATH.parent
    final_model_dir = config.FINAL_MODEL_PATH.parent
    
    if os.path.exists(initial_model_dir):
        print(f"自动清理旧的初始模型目录: {initial_model_dir}")
        shutil.rmtree(initial_model_dir)
    
    if os.path.exists(final_model_dir):
        print(f"自动清理旧的最终模型目录: {final_model_dir}")
        shutil.rmtree(final_model_dir)

    # 重新创建空的模型目录，以便保存新模型
    os.makedirs(initial_model_dir, exist_ok=True)
    os.makedirs(final_model_dir, exist_ok=True)


    # --- 步骤 1: 准备共享资源 (所有实验都需要) ---
    positive_data, negative_data = data_loader.get_and_prepare_cifar10()
    p_dataset, u_dataset, u_labels = data_loader.create_pu_dataset(positive_data, negative_data)
    p_embeddings, u_embeddings = feature_extractor.get_or_create_embeddings(p_dataset, u_dataset)
    
    all_results = {} # 创建一个字典来存储所有实验的结果

    # --- 步骤 2: 根据config开关，调度执行不同的实验 ---

    # --- 运行我们自己的核心算法 ---
    if config.RUN_OUR_METHOD:
        # 加载或训练初始模型
        if os.path.exists(config.INITIAL_MODEL_PATH) and not config.RUN_INITIAL_MODEL_ONLY:
             print(f"\n发现已存在的初始模型，直接加载: {config.INITIAL_MODEL_PATH}")
             initial_model = SimpleMLP(input_dim=config.EMBEDDING_DIM, hidden_dims=config.HIDDEN_DIMS)
             initial_model.load_state_dict(torch.load(config.INITIAL_MODEL_PATH))
        else:
            initial_model = initial_trainer.train_initial_model(p_embeddings, u_embeddings)
        
        if config.RUN_INITIAL_MODEL_ONLY:
            all_results['Initial Model (Phase 1)'] = evaluator.evaluate(initial_model, u_embeddings, u_labels, "Initial Model (Phase 1)")

        # 加载或训练最终模型
        if os.path.exists(config.FINAL_MODEL_PATH):
            print(f"\n发现已存在的最终模型，直接加载: {config.FINAL_MODEL_PATH}")
            final_model = SimpleMLP(input_dim=config.EMBEDDING_DIM, hidden_dims=config.HIDDEN_DIMS)
            final_model.load_state_dict(torch.load(config.FINAL_MODEL_PATH))
        else:
            final_model = em_trainer.train_final_model_with_em(initial_model, p_embeddings, u_embeddings)
        
        all_results['Our EM Model (Final)'] = evaluator.evaluate(final_model, u_embeddings, u_labels, "Our EM Model (Final)")

    # --- 运行基准实验 ---
    if config.RUN_FULLY_SUPERVISED:
        oracle_model = baselines.run_fully_supervised(p_embeddings, u_embeddings, u_labels)
        all_results['Fully Supervised (Oracle)'] = evaluator.evaluate(oracle_model, u_embeddings, u_labels, "Fully Supervised (Oracle)")

    if config.RUN_ZERO_SHOT:
        # 零样本需要原始图像
        u_images_dataset = TensorDataset(torch.stack([d[0] for d in positive_data[config.N_POSITIVE:]] + [d[0] for d in negative_data]))
        predictions, true_labels_np = baselines.run_zero_shot(u_images_dataset, u_labels)
        all_results['Zero-Shot'] = evaluator.evaluate_zero_shot(predictions, true_labels_np)

    if config.RUN_NAIVE_SAMPLING:
        naive_model = baselines.run_naive_sampling(p_embeddings, u_embeddings)
        all_results['Naive Sampling'] = evaluator.evaluate(naive_model, u_embeddings, u_labels, "Naive Sampling")

    # --- 步骤 3: 汇总和展示所有结果 ---
    evaluator.save_summary_table(all_results)
    
    print("="*60)
    print("所有流程执行完毕。")
    print("="*60)


if __name__ == '__main__':
    main()